import streamlit as st
import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
from sklearn.datasets import make_classification
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
import time
import plotly.graph_objects as go
from plotly.subplots import make_subplots


def show():
    st.markdown('<h2 class="sub-header">优化器比较</h2>', unsafe_allow_html=True)

    st.markdown("""
    <div class="content">
        <p>本模块比较不同优化器在神经网络训练中的性能表现。</p>
        <p>优化器是深度学习中的关键组件，负责更新网络参数以最小化损失函数。</p>
    </div>
    """, unsafe_allow_html=True)

    # 优化器对比表格
    st.subheader("优化器对比")

    st.markdown("""
    | 优化器 | 核心思想 | 优点 | 缺点 | 适用场景 |
    |--------|---------|------|------|---------|
    | GD | 使用全部数据计算梯度 | 收敛稳定 | 计算慢，内存要求高 | 小数据集 |
    | SGD | 使用单个样本更新 | 计算快，可在线学习 | 波动大，收敛不稳定 | 大规模数据 |
    | MBGD | 使用小批量样本 | 平衡速度与稳定性 | 需要调整批量大小 | 大多数深度学习任务 |
    | Momentum | 引入动量项 | 加速收敛，减少震荡 | 需要调整动量参数 | 高维、稀疏梯度 |
    | AdaGrad | 自适应学习率 | 稀疏数据表现好（NLP） | 学习率持续下降 | 自然语言处理 |
    | RMSProp | 指数加权平均 | 解决AdaGrad学习率问题 | 需要调整衰减率 | 非平稳目标函数 |
    | Adam | 结合Momentum和RMSProp | 自适应学习率，收敛快 | 超参数较多 | 大多数深度学习应用 |
    """)

    # 选择优化器
    st.subheader("选择优化器进行比较")

    optimizer_options = {
        "SGD": optim.SGD,
        "Momentum": lambda params, lr: optim.SGD(params, lr=lr, momentum=0.9),
        "AdaGrad": optim.Adagrad,
        "RMSProp": optim.RMSprop,
        "Adam": optim.Adam
    }

    selected_optimizers = st.multiselect(
        "选择要比较的优化器",
        list(optimizer_options.keys()),
        default=["SGD", "Adam"]
    )

    # 训练参数
    st.subheader("训练参数")

    col1, col2 = st.columns(2)
    with col1:
        learning_rate = st.slider("学习率", 0.0001, 0.1, 0.01, 0.0001, format="%.4f")
        epochs = st.slider("训练轮数", 10, 500, 100, 10)
    with col2:
        hidden_units = st.slider("隐藏层神经元数量", 16, 256, 64, 16)
        batch_size = st.slider("批处理大小", 16, 256, 64, 16)

    # 生成数据
    X, y = make_classification(
        n_samples=1000, n_features=20, n_redundant=2,
        n_informative=10, random_state=42, n_clusters_per_class=1
    )

    # 数据标准化和分割
    scaler = StandardScaler()
    X = scaler.fit_transform(X)
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

    # 转换为PyTorch张量
    X_train_tensor = torch.FloatTensor(X_train)
    y_train_tensor = torch.LongTensor(y_train)
    X_test_tensor = torch.FloatTensor(X_test)
    y_test_tensor = torch.LongTensor(y_test)

    # 创建数据加载器
    train_dataset = TensorDataset(X_train_tensor, y_train_tensor)
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

    # 定义神经网络模型
    class SimpleNN(nn.Module):
        def __init__(self, input_size, hidden_units, num_classes):
            super(SimpleNN, self).__init__()
            self.fc1 = nn.Linear(input_size, hidden_units)
            self.relu = nn.ReLU()
            self.fc2 = nn.Linear(hidden_units, num_classes)

        def forward(self, x):
            out = self.fc1(x)
            out = self.relu(out)
            out = self.fc2(out)
            return out

    # 训练按钮
    if st.button("开始比较优化器"):
        if not selected_optimizers:
            st.warning("请至少选择一个优化器")
            return

        # 准备结果存储
        results = {}

        # 对每个选中的优化器进行训练
        for opt_name in selected_optimizers:
            st.write(f"训练 {opt_name} 优化器...")

            # 初始化模型
            model = SimpleNN(X.shape[1], hidden_units, 2)
            criterion = nn.CrossEntropyLoss()

            # 创建优化器
            if opt_name == "Momentum":
                optimizer = optimizer_options[opt_name](model.parameters(), lr=learning_rate)
            else:
                optimizer = optimizer_options[opt_name](model.parameters(), lr=learning_rate)

            # 训练模型
            losses = []
            accuracies = []

            start_time = time.time()

            for epoch in range(epochs):
                model.train()
                epoch_loss = 0
                correct = 0
                total = 0

                for batch_X, batch_y in train_loader:
                    # 前向传播
                    outputs = model(batch_X)
                    loss = criterion(outputs, batch_y)

                    # 反向传播和优化
                    optimizer.zero_grad()
                    loss.backward()
                    optimizer.step()

                    # 计算准确率
                    _, predicted = torch.max(outputs.data, 1)
                    total += batch_y.size(0)
                    correct += (predicted == batch_y).sum().item()

                    epoch_loss += loss.item()

                # 记录损失和准确率
                avg_loss = epoch_loss / len(train_loader)
                accuracy = 100 * correct / total

                losses.append(avg_loss)
                accuracies.append(accuracy)

            training_time = time.time() - start_time

            # 评估模型
            model.eval()
            with torch.no_grad():
                test_outputs = model(X_test_tensor)
                _, test_predicted = torch.max(test_outputs.data, 1)
                test_accuracy = 100 * (test_predicted == y_test_tensor).sum().item() / y_test_tensor.size(0)

            # 存储结果
            results[opt_name] = {
                "losses": losses,
                "accuracies": accuracies,
                "training_time": training_time,
                "test_accuracy": test_accuracy
            }

        # 可视化结果
        st.subheader("优化器性能比较")

        # 创建子图
        fig = make_subplots(
            rows=2, cols=2,
            subplot_titles=("训练损失", "训练准确率", "测试准确率", "训练时间"),
            specs=[[{"type": "scatter"}, {"type": "scatter"}],
                   [{"type": "bar"}, {"type": "bar"}]]
        )

        # 添加训练损失曲线
        for opt_name, result in results.items():
            fig.add_trace(
                go.Scatter(
                    x=list(range(1, epochs + 1)),
                    y=result["losses"],
                    mode="lines",
                    name=f"{opt_name} 损失",
                    legendgroup=opt_name
                ),
                row=1, col=1
            )

        # 添加训练准确率曲线
        for opt_name, result in results.items():
            fig.add_trace(
                go.Scatter(
                    x=list(range(1, epochs + 1)),
                    y=result["accuracies"],
                    mode="lines",
                    name=f"{opt_name} 准确率",
                    legendgroup=opt_name,
                    showlegend=False
                ),
                row=1, col=2
            )

        # 添加测试准确率柱状图
        test_accuracies = [result["test_accuracy"] for result in results.values()]
        fig.add_trace(
            go.Bar(
                x=list(results.keys()),
                y=test_accuracies,
                name="测试准确率",
                marker_color="blue"
            ),
            row=2, col=1
        )

        # 添加训练时间柱状图
        training_times = [result["training_time"] for result in results.values()]
        fig.add_trace(
            go.Bar(
                x=list(results.keys()),
                y=training_times,
                name="训练时间",
                marker_color="green"
            ),
            row=2, col=2
        )

        # 更新布局
        fig.update_layout(
            height=800,
            title_text="优化器性能比较",
            showlegend=True
        )

        fig.update_xaxes(title_text="Epoch", row=1, col=1)
        fig.update_xaxes(title_text="Epoch", row=1, col=2)
        fig.update_xaxes(title_text="优化器", row=2, col=1)
        fig.update_xaxes(title_text="优化器", row=2, col=2)

        fig.update_yaxes(title_text="损失", row=1, col=1)
        fig.update_yaxes(title_text="准确率 (%)", row=1, col=2)
        fig.update_yaxes(title_text="准确率 (%)", row=2, col=1)
        fig.update_yaxes(title_text="时间 (秒)", row=2, col=2)

        st.plotly_chart(fig, use_container_width=True)

        # 显示详细结果
        st.subheader("详细结果")

        for opt_name, result in results.items():
            st.write(f"**{opt_name}**:")
            st.write(f"- 最终训练损失: {result['losses'][-1]:.4f}")
            st.write(f"- 最终训练准确率: {result['accuracies'][-1]:.2f}%")
            st.write(f"- 测试准确率: {result['test_accuracy']:.2f}%")
            st.write(f"- 训练时间: {result['training_time']:.2f} 秒")
            st.write("---")

    # 优化器原理说明
    st.subheader("优化器原理")

    with st.expander("SGD (随机梯度下降)"):
        st.markdown("""
        SGD使用单个样本或小批量样本计算梯度并更新参数：

        $$\\theta_{t+1} = \\theta_t - \\eta \\nabla J(\\theta_t; x^{(i)}, y^{(i)})$$

        其中：
        - $\\theta_t$ 是第t次迭代的参数
        - $\\eta$ 是学习率
        - $\\nabla J$ 是损失函数关于参数的梯度
        - $(x^{(i)}, y^{(i)})$ 是训练样本
        """)

    with st.expander("Momentum (动量)"):
        st.markdown("""
        Momentum引入了动量项，加速SGD在相关方向的收敛：

        $$v_t = \\gamma v_{t-1} + \\eta \\nabla J(\\theta_t)$$
        $$\\theta_{t+1} = \\theta_t - v_t$$

        其中：
        - $v_t$ 是动量项
        - $\\gamma$ 是动量系数（通常设为0.9）
        """)

    with st.expander("AdaGrad (自适应梯度)"):
        st.markdown("""
        AdaGrad为每个参数自适应地调整学习率：

        $$G_t = G_{t-1} + (\\nabla J(\\theta_t))^2$$
        $$\\theta_{t+1} = \\theta_t - \\frac{\\eta}{\\sqrt{G_t + \\epsilon}} \\odot \\nabla J(\\theta_t)$$

        其中：
        - $G_t$ 是梯度平方的累积
        - $\\epsilon$ 是小常数，防止除零
        - $\\odot$ 表示逐元素乘法
        """)

    with st.expander("RMSProp"):
        st.markdown("""
        RMSProp使用指数加权移动平均解决AdaGrad学习率持续下降的问题：

        $$E[g^2]_t = \\gamma E[g^2]_{t-1} + (1-\\gamma)(\\nabla J(\\theta_t))^2$$
        $$\\theta_{t+1} = \\theta_t - \\frac{\\eta}{\\sqrt{E[g^2]_t + \\epsilon}} \\nabla J(\\theta_t)$$

        其中：
        - $E[g^2]_t$ 是梯度平方的指数加权移动平均
        - $\\gamma$ 是衰减率（通常设为0.9）
        """)

    with st.expander("Adam (自适应矩估计)"):
        st.markdown("""
        Adam结合了Momentum和RMSProp的优点：

        $$m_t = \\beta_1 m_{t-1} + (1-\\beta_1) \\nabla J(\\theta_t)$$
        $$v_t = \\beta_2 v_{t-1} + (1-\\beta_2) (\\nabla J(\\theta_t))^2$$
        $$\\hat{m}_t = \\frac{m_t}{1-\\beta_1^t}$$
        $$\\hat{v}_t = \\frac{v_t}{1-\\beta_2^t}$$
        $$\\theta_{t+1} = \\theta_t - \\frac{\\eta}{\\sqrt{\\hat{v}_t} + \\epsilon} \\hat{m}_t$$

        其中：
        - $m_t$ 和 $v_t$ 分别是梯度的一阶矩和二阶矩估计
        - $\\beta_1$ 和 $\\beta_2$ 是衰减率（通常设为0.9和0.999）
        - $\\hat{m}_t$ 和 $\\hat{v}_t$ 是偏差校正后的估计
        """)

    # 优化器选择指南
    st.subheader("优化器选择指南")

    st.markdown("""
    1. **对于大多数情况**：Adam通常是首选，因为它结合了Momentum和RMSProp的优点
    2. **对于稀疏数据**：AdaGrad、RMSProp或Adam更适合
    3. **对于需要快速收敛**：使用Momentum或Adam
    4. **对于简单问题**：SGD可能足够，并且更容易调整
    5. **对于需要稳定训练**：使用SGD with Momentum
    6. **对于需要精细调整**：SGD通常提供更好的最终性能，但需要更多超参数调整
    """)

    # 代码示例
    st.subheader("PyTorch优化器使用示例")

    st.code("""
import torch
import torch.nn as nn
import torch.optim as optim

# 定义模型
model = nn.Linear(10, 2)

# 定义损失函数
criterion = nn.CrossEntropyLoss()

# 使用不同的优化器
optimizer_sgd = optim.SGD(model.parameters(), lr=0.01)
optimizer_momentum = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
optimizer_adagrad = optim.Adagrad(model.parameters(), lr=0.01)
optimizer_rmsprop = optim.RMSprop(model.parameters(), lr=0.01)
optimizer_adam = optim.Adam(model.parameters(), lr=0.01)

# 训练循环
for epoch in range(100):
    # 前向传播
    outputs = model(inputs)
    loss = criterion(outputs, labels)

    # 反向传播和优化
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    # 打印训练信息
    if (epoch+1) % 10 == 0:
        print(f'Epoch [{epoch+1}/100], Loss: {loss.item():.4f}')
    """, language="python")


if __name__ == "__main__":
    show()